package top.atcra.rpc.protection.loadBalancer;

import cn.hutool.core.lang.Assert;
import lombok.extern.slf4j.Slf4j;
import top.atcra.rpc.discovery.register.ServiceRegistry;
import top.atcra.rpc.protection.health.ServiceHealthDetector;
import top.atcra.rpc.protection.loadBalancer.selector.ServiceSelector;

import java.net.InetSocketAddress;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.Function;

@Slf4j
public class LoadBalancer {
    protected final ServiceRegistry serviceRegistry;
    protected ServiceHealthDetector serviceHealthDetector;

    public LoadBalancer enableHealthCheck(ServiceHealthDetector serviceHealthDetector) {
        this.serviceHealthDetector = serviceHealthDetector;
        return this;
    }

    public LoadBalancer(ServiceRegistry registry, Function<List<InetSocketAddress>, ServiceSelector> getSelector) {
        this.serviceRegistry = registry;
        this.getSelector = getSelector;
    }

    // 服务名→服务选择器，选择器内部缓存了可用的服务列表
    protected final Map<String, ServiceSelector> cache = new ConcurrentHashMap<>(8);
    // 根据服务列表获取一个选择器的函数
    protected final Function<List<InetSocketAddress>, ServiceSelector> getSelector;

    public ServiceSelector initSelector(String serviceName, String group) {
        var cacheKey = serviceName + group;
        var serviceList = serviceRegistry.lookup(serviceName, group);
        Assert.notEmpty(serviceList, serviceName + "的可用服务列表为空");
        if (serviceHealthDetector != null) {
            // 检查服务是否可用
            serviceHealthDetector.checkHealth(serviceList);
        }
        ServiceSelector selector = getSelector.apply(serviceList);
        log.info("服务{}({})的可用服务结点有:{}", serviceName, group, serviceList);
        // 将select放入缓存当中
        cache.put(cacheKey, selector);
//        if (log.isDebugEnabled()) {
//            log.debug("为服务{}({})创建了一个选择器【{}】", serviceName,group, selector.getClass().getName());
//        }

        return selector;
    }

    public InetSocketAddress getAvailableService(String serviceName, String group, Object key) {
        var cacheKey = serviceName + group;
        var selector = cache.get(cacheKey);
        if (selector == null) {
            selector = initSelector(serviceName, group);
        }
        // 获取可用节点
        return selector.getAddress(key);
    }

    public void reBalance(String serviceName, String group, List<InetSocketAddress> addresses) {
        cache.put(serviceName + group, getSelector.apply(addresses));
    }
}
